import torch # PyTorch的核心库，用于张量运算和自动求导。
import torch.nn as nn # 包含各种神经网络层、激活函数、损失函数等模块。
import torch.optim as optim # 优化器模块，用于更新模型参数。
import torchvision # 提供常用的数据集、数据预处理和图像转换操作
import torchvision.transforms as transforms #
from torch.utils.data import DataLoader # 用来将数据集切分成一个个小批量，便于训练时逐批处理。
'''从网络结构、数据增强、批归一化（Batch Normalization）以及学习率调度器等多个方面入手，尽可能提高MNIST手写数字识别的准确率。
数据增强
具体措施：在训练数据上加入了随机旋转（RandomRotation），使得训练图像在一定范围内发生旋转变化。
原理：数据增强可以扩充训练样本的多样性，帮助模型学习到更具泛化能力的特征，降低过拟合风险，从而提高测试时的准确率。

网络结构的改进
增加卷积层：在原有两层卷积的基础上，增加了一层卷积层，输出通道数增加到128。
原理：更深的网络能够提取更丰富、更抽象的特征信息，使得模型对图像细节和变化更敏感，从而提升分类性能。
增大全连接层神经元数量：将全连接层的神经元数量从128扩展到256。
原理：增加神经元数量可以提升模型的表达能力，捕捉更复杂的非线性关系，帮助分类器更精确地区分不同数字。

批归一化（Batch Normalization）
具体措施：在每个卷积层之后加入了批归一化层。
原理：批归一化可以稳定每一层输入的分布，加速网络收敛，同时也具有一定的正则化作用，能减少内部协变量偏移，从而提高模型的鲁棒性和准确率。

学习率调度器（Learning Rate Scheduler）
具体措施：使用了 StepLR 调度器，在固定轮次后降低学习率。
原理：初期使用较高的学习率可以快速找到较好的解空间，后期降低学习率则有助于精细调整模型参数，使模型更稳定地收敛到更优解，从而提升最终的识别效果。

增加训练轮数
具体措施：将训练轮数从10轮增至20轮。
原理：更多的训练轮次能让模型充分学习和调整参数，从而获得更好的拟合效果，达到更高的准确率。

实验效果：测试集上的准确率为从98.85%提升到了99.57%

'''
# 检查是否有可用的GPU，如果有则使用GPU加速
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# ----------------------- 超参数设置 -----------------------
batch_size = 64  # 每个mini-batch的样本数（每次训练处理的数据数量）
learning_rate = 0.001  # 初始学习率（优化器更新参数时的步长，决定学习速度。）
num_epochs = 20  # 训练轮数，可以适当增多
step_size = 10  # 学习率调度器参数，每step_size轮降低一次学习率 （用于学习率调度器，控制训练过程中降低学习率的策略。）
gamma = 0.5  # 学习率衰减因子

# ----------------------- 数据预处理和数据增强 -----------------------
# 为了增强模型的泛化能力，对训练集进行数据增强（例如随机旋转）
train_transform = transforms.Compose([
    transforms.RandomRotation(10),  # 随机将图像旋转±10度范围内，模拟真实环境下手写数字的不同倾斜情况，从而使模型具备更好的泛化能力
    transforms.ToTensor(), # 将图片转换为PyTorch张量，便于后续运算。
    transforms.Normalize((0.1307,), (0.3081,)) #  对图像数据进行归一化处理，使图像的像素值分布接近标准正态分布（均值0.1307，标准差0.3081），这有助于模型训练稳定。
])
# 测试集仅进行归一化处理
test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 下载MNIST数据集
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=train_transform)
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=test_transform)

# DataLoader加载数据集
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)


# ----------------------- 定义改进的卷积神经网络 -----------------------
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 第一层卷积：输入1通道，输出32通道；卷积核大小3，padding=1保持尺寸
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(32)  # 批归一化，有助于训练稳定 在每个卷积层后加入，帮助稳定训练过程和加速收敛。

        # 第二层卷积：32->64
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(64)

        # 第一处池化层：将尺寸减半（28x28 -> 14x14） 通过取区域内的最大值来降低特征图尺寸，减少计算量和参数量，同时捕获最重要的特征。
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        # 第三层卷积：64->128
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
        self.bn3 = nn.BatchNorm2d(128)

        # 第二处池化层：继续减半（14x14 -> 7x7）
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        # Dropout层：防止过拟合 在训练过程中随机“丢弃”部分神经元，防止模型过拟合。
        self.dropout = nn.Dropout(0.25)

        # 全连接层：输入特征数 = 128通道 * 7 * 7 = 6272 5.4.4.4.将前面提取到的特征映射到最终的分类结果（10个类别对应数字0~9）。
        self.fc1 = nn.Linear(in_features=128 * 7 * 7, out_features=256)
        # 输出层：256 -> 10（对应数字0-9）
        self.fc2 = nn.Linear(in_features=256, out_features=10)

    def forward(self, x):
        # x 代表输入数据。在这个例子中，x 是一个包含多张手写数字图像的批次（batch）。每张图像的尺寸是 28x28 像素，且是单通道（灰度图像）。
        # 第一层卷积 -> 批归一化 -> ReLU激活
        x = torch.relu(self.bn1(self.conv1(x))) #将输入图像通过第一层卷积操作，提取出特征图（从1通道变成32通道）。然后归一化
        # 第二层卷积 -> 批归一化 -> ReLU激活
        x = torch.relu(self.bn2(self.conv2(x)))
        # 第一处池化
        x = self.pool(x)
        # 第三层卷积 -> 批归一化 -> ReLU激活
        x = torch.relu(self.bn3(self.conv3(x)))
        # 第二处池化
        x = self.pool2(x)
        x = self.dropout(x)
        # 展平为二维张量，尺寸变为(batch_size, 128*7*7)
        x = x.view(x.size(0), -1) #将多维的特征图展平为一维。这里 x.size(0) 是批次的大小（batch size），-1 会自动计算出剩余维度的大小。这样做是为了将提取到的特征输入到全连接层。
        # 全连接层1 -> ReLU激活
        x = torch.relu(self.fc1(x))
        x = self.dropout(x)
        # 输出层
        x = self.fc2(x)
        return x


# 将模型实例化并移动到设备上
model = CNN().to(device)

# ----------------------- 定义损失函数和优化器 -----------------------
criterion = nn.CrossEntropyLoss()  # 多分类交叉熵损失，使用交叉熵损失函数（CrossEntropyLoss），适合多分类问题。这个函数会自动将输出得分转换为概率分布，并计算预测概率与真实标签之间的误差。
optimizer = optim.Adam(model.parameters(), lr=learning_rate) # 采用Adam优化器，它是一种自适应学习率算法，能够根据梯度的变化自动调整学习率，加快训练收敛速度。
# 使用学习率调度器，每step_size轮降低学习率,使用 StepLR 调度器，在每隔一定训练轮数后（step_size），将学习率乘以一个衰减因子（gamma）。这种策略在训练后期降低学习率，有助于模型细致调整参数，进一步提升性能。
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)

# ----------------------- 模型训练 -----------------------
# 进入训练状态，价值数据集，小连招，调整学习率
print("开始训练...")
for epoch in range(num_epochs):
    model.train()  # 设置为训练模式 使用 model.train() 告诉模型当前处于训练状态（这时Dropout等层会生效）。
    running_loss = 0.0
    for i, (images, labels) in enumerate(train_loader):
        images, labels = images.to(device), labels.to(device)
        # 使用 train_loader 将数据集分批加载，每个批次都进行一次前向传播、计算损失、反向传播以及参数更新。
        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播及参数更新
        optimizer.zero_grad() #清空上一步的梯度，防止梯度累积
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        # 每训练一定数量的mini-batch输出一次当前平均损失，方便观察训练进度
        if (i + 1) % 100 == 0:
            print(
                f'轮数 [{epoch + 1}/{num_epochs}], 第 [{i + 1}/{len(train_loader)}] 个mini-batch, 平均损失: {running_loss / 100:.4f}')
            running_loss = 0.0
    # 每个epoch结束后调整学习率
    scheduler.step()

# ----------------------- 模型测试 -----------------------
# 模型测试：关闭梯度计算，设置为评估模式、遍历测试集，测试及统计结果
model.eval()  # 设置为评估模式，关闭Dropout等训练专用层
correct = 0
total = 0

with torch.no_grad():
    for images, labels in test_loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        # 获取预测结果（沿类别维度取最大值）
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'测试集上的准确率为: {100 * correct / total:.2f}%')
